#include <stdio.h>
#include "sep.h"
#include "septools.h"

typedef struct {
  int nmols;
  int nmols_A, nmols_B;
  double **posA, **posB;
} type_info;


void get_number_of_types(const char file[], double *lbox,
			 char typeA, char typeB, type_info *info){
  double tnow;
  char line[2048];

  FILE *fin = fopen(file, "r");
  if ( fin == NULL )
    sep_error("Couldn't read input file- BAILING OUT");

  fscanf(fin, "%lf %lf %lf %lf\n", &tnow, &lbox[0], &lbox[1], &lbox[2]);

  info->nmols = 0;
  info->nmols_A = 0;
  info->nmols_B = 0;

  while ( fgets(line, 2048, fin) != NULL ) {
    
    (info->nmols)++;

    if ( line[0] == typeA )
      (info->nmols_A)++;
    
    if ( line[0] == typeB ) // Allows for A-A
      (info->nmols_B)++;
  }

  fclose(fin);

  printf("Found %d molecules of type %c ", info->nmols_A, typeA);
  printf("Found %d molecules of type %c \n", info->nmols_B, typeB);

  fflush(stdout);

}


void read_entry_from_file(double **posA, double **posB, char typeA, 
			  char typeB, const char *file, type_info tinfo){
  double tnow, lx, ly, lz;
  char type;
  int nuau;
  double m, x[3], v[3], s[3], w[3], inertia[6]; // xx,yy,zz,xy,xz,yz component

 
  FILE *fin = fopen(file, "r");
  if ( fin == NULL )
    sep_error("Couldn't read input file- BAILING OUT");
  
  fscanf(fin, "%lf %lf %lf %lf\n", &tnow, &lx, &ly, &lz);

  int cA=0; int cB=0;
  for ( int i=0; i<tinfo.nmols; i++){

    int nread = fscanf(fin, "%c %d %lf %lf %lf %lf %lf %lf %lf %lf %lf %lf %lf %lf %lf %lf %lf %lf %lf %lf %lf\n", 
		       &type, &nuau, &m, 
		       &x[0], &x[1], &x[2], 
		       &v[0], &v[1], &v[2],
		       &s[0], &s[1], &s[2],
		       &w[0], &w[1], &w[2],
		       &inertia[0], &inertia[1], &inertia[2],
		       &inertia[3], &inertia[4], &inertia[5]);
    if ( nread != 21 )
      sep_error("Read failure - BAILING OUT");

    if ( type == typeA ){
      if ( cA >= tinfo.nmols_A ){
	fprintf(stderr, "Did you clean directory before running again?");
	sep_error("Index larger than number of molecules - BAILING OUT");
      }
      for ( int k=0; k<3; k++) 
	posA[cA][k] = x[k];
      cA ++;
    }
    if ( type == typeB ){
      if ( cB >= tinfo.nmols_B ){
	fprintf(stderr, "Did you clean directory before running again?");
	sep_error("Index larger than number of molecules - BAILING OUT");
      }
      for ( int k=0; k<3; k++) 
	posB[cB][k] = x[k];
      cB++;
    } 
  }

  fclose(fin);

} 

void rdfsample(double *hist, double **posA, double **posB, double lbox[3], 
	       double dg, int lvec, type_info tinfo){
  double r[3];

  for ( int i=0; i<tinfo.nmols_A; i++ ){
    for ( int j=0; j<tinfo.nmols_B; j++ ){
      double r2 = 0.0;
      for ( int k=0; k<3; k++){
        r[k]  = posA[i][k]-posB[j][k];
        sep_Wrap( r[k], lbox[k] );
        r2   += r[k]*r[k];
      }
        
      int index = (int)(sqrt(r2)/dg);
      if  (index < lvec ) hist[index] += 1.0;
    }
  }


}


int main(int argc, char **argv){
  type_info tinfo;
  double lbox[3], dg;

  if ( argc != 4 )
    sep_error("Arguments <typeA> <typeB> <length of vector>");

  char typeA = argv[1][0];
  char typeB = argv[2][0];
  int lvec = atoi(argv[3]);

  get_number_of_types("00000_mol.conf", lbox, typeA, typeB, &tinfo);

  if ( (lbox[0] > lbox[1]) && (lbox[1] > lbox[2]) )
    dg = lbox[2]/(2*lvec);
  if ( (lbox[0] > lbox[2]) && (lbox[2] > lbox[1]) )
    dg = lbox[1]/(2*lvec);
  else
    dg = lbox[0]/(2*lvec);

  double **posA = sep_matrix(tinfo.nmols_A, 3);
  double **posB = sep_matrix(tinfo.nmols_B, 3);
  double *hist = sep_vector(lvec);

  int file_counter = 0;
  while (1) {

    printf("\r");
    printf("Processing sample %d  ", file_counter);
    fflush(stdout);

    char file[256];
    sprintf(file, "%05d_mol.conf", file_counter);
    file_counter++;
    if ( file_exists(file) == 1 ){
      read_entry_from_file(posA, posB, typeA, typeB, file, tinfo);
      rdfsample(hist, posA, posB, lbox, dg, lvec, tinfo);
    }
    else {
      break;
    }
  }
 
  printf("\n");


  FILE *fout = fopen("mrdf.dat", "w");
  if ( fout == NULL )
    sep_error("Couldn't open file - BAILING OUT");

  // Remove self-self distance for 
  if ( typeA == typeB ) hist[0] = 0;

  int nmolsAB = tinfo.nmols_A+tinfo.nmols_B;
  double vol  = lbox[0]*lbox[1]*lbox[2];
  double dens = nmolsAB/vol;

  for ( int k=0; k<lvec; k++ ){
    double r   = dg*(k+0.5);
    
    double vi = pow(k*dg, 3.0);
    double vii = pow((k+1)*dg, 3.0);
    double nid = 4./3.*SEP_PI*(vii-vi)*dens;

    double fac = 1.0/(0.25*file_counter*nmolsAB*nid);
    fprintf(fout, "%f %f\n", r, hist[k]*fac);
  }

  fclose(fout);

  sep_free_matrix(posA, tinfo.nmols_A);
  sep_free_matrix(posB, tinfo.nmols_B);
  free(hist);

  printf("Done - wrote result in mrdf.dat\n");

  return 0;
}

